# gpu调度分两个阶段：
# 1. 当能够不跨leaf通信时，这一阶段不涉及connection manager
# 2. 当需要跨leaf通信时，此时有两种情况：

import numpy as np
from ortools.linear_solver import pywraplp
from ortools.graph.python import linear_sum_assignment


class ConnectionManager:
    def __init__(self, gpu_num = 512, server_num = 64, leaf_num = 16, spine_num = 16):
        self.gpu_num = gpu_num
        self.server_num = server_num
        self.leaf_num = leaf_num
        self.spine_num = spine_num
        self.server_per_leaf = int(server_num/leaf_num)
        self.gpu_per_server = int(gpu_num/server_num)

        self.leaf_to_spine_remain_port_num = {}
        for leaf_id in range(leaf_num):
            for to_spine_id in range(spine_num):
                if leaf_id not in self.leaf_to_spine_remain_port_num:
                    self.leaf_to_spine_remain_port_num[leaf_id] = {}
                self.leaf_to_spine_remain_port_num[leaf_id][to_spine_id] = int(self.server_per_leaf*self.gpu_per_server/self.spine_num)
        self.leaf_remain_port_map = {}
        for leaf_id in range(leaf_num):
            self.leaf_remain_port_map[leaf_id] = int(self.gpu_num/self.leaf_num)
        self.spine_remain_port_map = {}
        for spine_id in range(spine_num):
            self.spine_remain_port_map[spine_id] = int(self.gpu_num/self.spine_num)



    def print_connection_info(self):
        for leaf_id in range(self.leaf_num):
            for spine_id in range(self.spine_num):
                # if self.leaf_to_spine_remain_port_num[leaf_id][spine_id] >0:
                print(leaf_id,spine_id,self.leaf_to_spine_remain_port_num[leaf_id][spine_id])

    def release_connection_resource(self, job_allocated_leaf_spine_link):
        self.check_valid_network()
        for leaf_id in job_allocated_leaf_spine_link:
            for spine_id in job_allocated_leaf_spine_link[leaf_id]:
                self.leaf_to_spine_remain_port_num[leaf_id][spine_id] += job_allocated_leaf_spine_link[leaf_id][spine_id]
                self.leaf_remain_port_map[leaf_id] += job_allocated_leaf_spine_link[leaf_id][spine_id]
                self.spine_remain_port_map[spine_id] += job_allocated_leaf_spine_link[leaf_id][spine_id]
                if self.leaf_to_spine_remain_port_num[leaf_id][spine_id]<0:
                    print(job_allocated_leaf_spine_link[leaf_id][spine_id])
                assert self.leaf_to_spine_remain_port_num[leaf_id][spine_id] >= 0
        self.check_valid_network()

    def check_valid_network(self):
        temp_spine_port_remain = {}
        for spine_id in range(self.spine_num):
            if spine_id not in temp_spine_port_remain:
                temp_spine_port_remain[spine_id] = 0 
        for spine_id in range(self.spine_num):
            for leaf_id in range(self.leaf_num):
                temp_spine_port_remain[spine_id] += self.leaf_to_spine_remain_port_num[leaf_id][spine_id]
        for spine_id in range(self.spine_num):
            if temp_spine_port_remain[spine_id] != self.spine_remain_port_map[spine_id]:
                print(self.leaf_num, self.spine_num)
                print(spine_id, temp_spine_port_remain[spine_id] ,self.spine_remain_port_map[spine_id])
            assert temp_spine_port_remain[spine_id] == self.spine_remain_port_map[spine_id]

    def find_valid_gpu_for_specific_spine(self, require_gpu_num, require_spine_id, server_remain_gpuNum_map,job_allocated_leaf_spine_link, used_spine_port_num_pair, leaf_remain_empt_server_list):
        solver = pywraplp.Solver('SolveIntegerProblem',
                             pywraplp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)
        
        name_list_x_i = []
        name_list_y_j = []
        name_list_z_i_j = []
        for i in range(self.server_num):
            name_list_x_i.append(str(i))
        for j in range(self.spine_num):
            name_list_y_j.append(str(j))
        for i in range(self.server_num):
            for j in range(self.spine_num):
                name_list_z_i_j.append(str(i)+'_'+str(j))
        y_j = {}
        for it in name_list_y_j:
            y_j[it] = solver.IntVar(0.0, self.spine_remain_port_map[j]*1.0, 'y_j'+str(it))
        z_i_j = {}
        for it in name_list_z_i_j:
            z_i_j[it] =solver.IntVar(0.0, 1.0*int(self.server_per_leaf*self.gpu_per_server/self.spine_num), 'z_i_j'+str(it))
        xnum_i = {}
        for it in name_list_x_i:
            xnum_i[it] = solver.IntVar(0.0, 1.0, 'xnum_i'+str(it))
        objective = solver.Objective()

        
        solver.Add(z_i_j["0_0"]+z_i_j["1_0"]+z_i_j["2_0"]+z_i_j["3_0"]+z_i_j["4_0"]+z_i_j["5_0"]+z_i_j["6_0"]+z_i_j["7_0"] <=  y_j[str(0)] )

        solver.Add(z_i_j["0_0"]+z_i_j["1_0"]+z_i_j["2_0"]+z_i_j["3_0"] <= self.leaf_to_spine_remain_port_num[0][0] )
        solver.Add(z_i_j["4_0"]+z_i_j["5_0"]+z_i_j["6_0"]+z_i_j["7_0"] <= self.leaf_to_spine_remain_port_num[1][0] )

        solver.Add(z_i_j["0_0"] <= server_remain_gpuNum_map[0] )
        solver.Add(z_i_j["1_0"] <= server_remain_gpuNum_map[1] )
        solver.Add(z_i_j["2_0"] <= server_remain_gpuNum_map[2] )
        solver.Add(z_i_j["3_0"] <= server_remain_gpuNum_map[3] )
        solver.Add(z_i_j["4_0"] <= server_remain_gpuNum_map[4] )
        solver.Add(z_i_j["5_0"] <= server_remain_gpuNum_map[5] )
        solver.Add(z_i_j["6_0"] <= server_remain_gpuNum_map[6] )
        solver.Add(z_i_j["7_0"] <= server_remain_gpuNum_map[7] )

        solver.Add(z_i_j["0_0"]+z_i_j["1_0"]+z_i_j["2_0"]+z_i_j["3_0"]+z_i_j["4_0"]+z_i_j["5_0"]+z_i_j["6_0"]+z_i_j["7_0"] == require_gpu_num)

        solver.Add(y_j["0"] == require_gpu_num)

        solver.Add(self.gpu_per_server*xnum_i[str(0)]==(z_i_j["0_0"]))
        solver.Add(self.gpu_per_server*xnum_i[str(1)]==(z_i_j["1_0"]))
        solver.Add(self.gpu_per_server*xnum_i[str(2)]==(z_i_j["2_0"]))
        solver.Add(self.gpu_per_server*xnum_i[str(3)]==(z_i_j["3_0"]))
        solver.Add(self.gpu_per_server*xnum_i[str(4)]==(z_i_j["4_0"]))
        solver.Add(self.gpu_per_server*xnum_i[str(5)]==(z_i_j["5_0"]))
        solver.Add(self.gpu_per_server*xnum_i[str(6)]==(z_i_j["6_0"]))
        solver.Add(self.gpu_per_server*xnum_i[str(7)]==(z_i_j["7_0"]))

        solver.Add(z_i_j["0_0"] <= y_j[str(0)] )
        solver.Add(z_i_j["1_0"] <= y_j[str(0)] )
        solver.Add(z_i_j["2_0"] <= y_j[str(0)] )
        solver.Add(z_i_j["3_0"] <= y_j[str(0)] )
        solver.Add(z_i_j["4_0"] <= y_j[str(0)] )
        solver.Add(z_i_j["5_0"] <= y_j[str(0)] )
        solver.Add(z_i_j["6_0"] <= y_j[str(0)] )
        solver.Add(z_i_j["7_0"] <= y_j[str(0)] )


        status = solver.Solve()
        if status == pywraplp.Solver.OPTIMAL:
            Z_i_j_solution = {}
            xnum_i_solution = {}
            for it in name_list_z_i_j:
                Z_i_j_solution[it] = z_i_j[it].solution_value()
            for it in name_list_x_i:
                xnum_i_solution[it] = xnum_i[it].solution_value()
            for it in name_list_z_i_j:
                divid_index = it.split("_")
                for id in range(len(divid_index)):
                    divid_index[id] = int(divid_index[id])
                if round(Z_i_j_solution[it]) >= 1:
                    to_use_link_num = round(Z_i_j_solution[it])
                    chosen_server_id = divid_index[0]
                    chosen_spine_id = divid_index[1]
                    chosen_leaf_id = int(chosen_server_id/self.server_per_leaf)
                    self.leaf_remain_port_map[chosen_leaf_id] -= to_use_link_num
                    self.spine_remain_port_map[chosen_spine_id] -= to_use_link_num
                    if chosen_leaf_id not in job_allocated_leaf_spine_link:
                        job_allocated_leaf_spine_link[chosen_leaf_id] = {}
                    if chosen_spine_id not in job_allocated_leaf_spine_link[chosen_leaf_id]:
                        job_allocated_leaf_spine_link[chosen_leaf_id][chosen_spine_id] = 0
                    job_allocated_leaf_spine_link[chosen_leaf_id][chosen_spine_id] += to_use_link_num
                    used_spine_port_num_pair[require_spine_id] = to_use_link_num
            server_occupy_gpuNum_map = {}
            for it in name_list_x_i:
                server_occupy_gpuNum_map[int(it)] = int(self.gpu_per_server*round(xnum_i_solution[it]))
            return True, server_occupy_gpuNum_map
        else:
            raise Exception("something wrong4 in solver")
            return False, None